from keras.models import Model
from keras.layers import Input, Embedding, Conv1D, GlobalMaxPool1D
from keras.layers import Dense, Dropout, BatchNormalization
import time
from keras.models import load_model
from evaluate import predict2both, predict2half, predict2top, f1_avg


def train(fact_train, labels_train, mode, num_words, max_len, kernel_size, dim, batch_size):
    print('start', time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))

    drop_out_rate = 0.2
    filter_nums = 512
    dense_size = 1000
    out_dim = 1
    loss = 'mae'
    activation = 'relu'
    if mode != 'imprisonment':
        out_dim = labels_train.shape[1]
        loss = 'binary_crossentropy'
        activation = 'sigmoid'

    data_input = Input(shape=[fact_train.shape[1]])
    word_vec = Embedding(input_dim=num_words + 1,
                         input_length=max_len,
                         output_dim=dim,
                         mask_zero=0,
                         name='Embedding')(data_input)
    x = word_vec
    x = Conv1D(filters=filter_nums, kernel_size=[kernel_size], strides=1, padding='same', activation='relu')(x)
    x = GlobalMaxPool1D()(x)
    x = BatchNormalization()(x)
    x = Dense(dense_size, activation="relu")(x)
    x = Dropout(drop_out_rate)(x)
    x = Dense(out_dim, activation=activation)(x)
    model = Model(inputs=data_input, outputs=x)
    model.compile(loss=loss,
                  optimizer='adam',
                  metrics=['accuracy'])
    # model.summary()

    model.fit(x=fact_train, y=labels_train, batch_size=batch_size, epochs=1, verbose=1)
    model.save('model/CNN_%s.h5' % mode)

    print('end', time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))


def test(mode, fact_test, labels_test):
    model = load_model('model/CNN_%s.h5' % mode)
    y = model.predict(fact_test[:])
    y1 = predict2top(y)
    y2 = predict2half(y)
    y3 = predict2both(y)

    print('accu:')
    # 只取最高置信度的准确率
    s1 = [(labels_test[i] == y1[i]).min() for i in range(len(y1))]
    print(sum(s1) / len(s1))
    # 只取置信度大于0.5的准确率
    s2 = [(labels_test[i] == y2[i]).min() for i in range(len(y1))]
    print(sum(s2) / len(s2))
    # 结合前两个
    s3 = [(labels_test[i] == y3[i]).min() for i in range(len(y1))]
    print(sum(s3) / len(s3))

    print('f1:')
    # 只取最高置信度的准确率
    s4 = f1_avg(y_pred=y1, y_true=labels_test)
    print(s4)
    # 只取置信度大于0.5的准确率
    s5 = f1_avg(y_pred=y2, y_true=labels_test)
    print(s5)
    # 结合前两个
    s6 = f1_avg(y_pred=y3, y_true=labels_test)
    print(s6)

